"""
Core image channel manipulation functions.
"""

from enum import Enum
from pathlib import Path
from typing import Tuple, Dict, Union
from PIL import Image, ImageChops

# Optional numpy import for some operations
try:
    import numpy as np
    NUMPY_AVAILABLE = True
except ImportError:
    NUMPY_AVAILABLE = False


class ImageType(Enum):
    """Enumeration of image types based on channel configuration."""
    UNKNOWN = "unknown"
    RGB = "rgb"  # 3-channel color
    RGBA = "rgba"  # 4-channel with alpha
    GRAYSCALE_1C = "grayscale_1c"  # Single channel grayscale (L mode)
    GRAYSCALE_3C = "grayscale_3c"  # 3-channel grayscale (RGB but all channels equal)
    PALETTE = "palette"  # Indexed color (P mode)
    CMYK = "cmyk"  # 4-channel CMYK
    LAB = "lab"  # 3-channel LAB
    HSV = "hsv"  # 3-channel HSV
    OTHER = "other"  # Other PIL modes


def load_image(image_input: Union[str, Path, Image.Image]) -> Image.Image:
    """
    Load an image from a file path or return the PIL Image object directly.

    Args:
        image_input: Either a file path (str or pathlib.Path) or a PIL Image object.

    Returns:
        PIL Image object.

    Raises:
        FileNotFoundError: If the file doesn't exist.
        ValueError: If the input is neither a string, Path, nor PIL Image.
    """
    if isinstance(image_input, (str, Path)):
        return Image.open(str(image_input))
    elif isinstance(image_input, Image.Image):
        return image_input
    else:
        raise ValueError(f"Unsupported input type: {type(image_input)}")


def detect_image_type(img: Image.Image) -> Dict[str, Union[ImageType, Tuple[int, ...], str]]:
    """
    Detect the type of an image based on its mode and channel content.

    Args:
        img: PIL Image object.

    Returns:
        Dictionary with keys:
        - 'type': ImageType enum
        - 'shape': (width, height, channels) tuple
        - 'mode': PIL mode string
        - 'channels': number of channels
        - 'is_actual_grayscale': bool (True if all channels have identical values)
    """
    mode = img.mode
    width, height = img.size

    # Map PIL modes to channel count
    mode_to_channels = {
        '1': 1,  # 1-bit pixels, black and white
        'L': 1,  # 8-bit pixels, black and white (grayscale)
        'P': 1,  # 8-bit pixels, mapped to any other mode using a color palette
        'RGB': 3,
        'RGBA': 4,
        'CMYK': 4,
        'YCbCr': 3,
        'LAB': 3,
        'HSV': 3,
        'I': 1,  # 32-bit signed integer pixels
        'F': 1,  # 32-bit floating point pixels
    }

    channels = mode_to_channels.get(mode, 1)

    # Determine base type
    image_type = ImageType.OTHER
    if mode == 'RGB':
        # Check if it's actually grayscale (all channels equal)
        if _is_grayscale_image(img):
            image_type = ImageType.GRAYSCALE_3C
        else:
            image_type = ImageType.RGB
    elif mode == 'RGBA':
        image_type = ImageType.RGBA
    elif mode == 'L':
        image_type = ImageType.GRAYSCALE_1C
    elif mode == 'P':
        image_type = ImageType.PALETTE
    elif mode == 'CMYK':
        image_type = ImageType.CMYK
    elif mode == 'LAB':
        image_type = ImageType.LAB
    elif mode == 'HSV':
        image_type = ImageType.HSV

    # Check if it's actually grayscale
    is_actual_grayscale = False
    if channels == 1:
        # Single channel images are grayscale by definition
        is_actual_grayscale = True
    elif channels >= 3:
        is_actual_grayscale = _is_grayscale_image(img)

    return {
        'type': image_type,
        'shape': (width, height, channels),
        'mode': mode,
        'channels': channels,
        'is_actual_grayscale': is_actual_grayscale,
    }


def _is_grayscale_image(img: Image.Image, threshold: int = 1) -> bool:
    """
    Check if an image is grayscale (all channels have identical values).

    Args:
        img: PIL Image object.
        threshold: Maximum allowed difference between channels (in 0-255 scale).

    Returns:
        True if all channels are identical within threshold.
    """
    mode = img.mode

    # Single channel images are grayscale by definition
    if mode in ('L', '1', 'I', 'F'):
        return True

    # For multi-channel images
    if mode == 'RGB':
        # Get individual channels
        r, g, b = img.split()
        # Check differences
        diff_rg = ImageChops.difference(r, g)
        diff_gb = ImageChops.difference(g, b)
        # Get maximum difference
        max_diff = ImageChops.difference(diff_rg, diff_gb)  # Not exactly max, but works
        # Actually compute max by checking each pixel
        # Simple implementation: check if any pixel difference > threshold
        # Convert to grayscale and check max value
        diff_rg_gray = diff_rg.convert('L')
        diff_gb_gray = diff_gb.convert('L')

        # Get histogram to find max value
        hist_rg = diff_rg_gray.histogram()
        hist_gb = diff_gb_gray.histogram()

        # Check if any difference > threshold
        max_val_rg = max(i for i, count in enumerate(hist_rg) if count > 0)
        max_val_gb = max(i for i, count in enumerate(hist_gb) if count > 0)

        return max_val_rg <= threshold and max_val_gb <= threshold

    elif mode == 'RGBA':
        # Check RGB channels, ignore alpha
        rgb_img = img.convert('RGB')
        return _is_grayscale_image(rgb_img, threshold)

    # For other multi-channel modes, assume not grayscale
    # (could be improved for CMYK, LAB, etc.)
    return False


def is_grayscale(img: Image.Image, threshold: int = 1) -> bool:
    """
    Check if an image is grayscale (single channel or all channels identical).

    Args:
        img: PIL Image object.
        threshold: Maximum allowed difference between channels (0-255).

    Returns:
        True if the image is grayscale.
    """
    return _is_grayscale_image(img, threshold)


def to_three_channel(img: Image.Image) -> Image.Image:
    """
    Convert an image to 3-channel format.

    - If single channel (L, 1, etc.): duplicate to 3 channels
    - If RGBA: remove alpha channel
    - If already 3-channel: return as-is
    - If palette: convert to RGB
    - Other modes will be converted to RGB

    Args:
        img: PIL Image object.

    Returns:
        3-channel RGB image.
    """
    mode = img.mode

    if mode == 'RGB':
        return img

    if mode == 'L' or mode == '1' or mode == 'I' or mode == 'F':
        # Single channel grayscale -> RGB
        return img.convert('RGB')

    if mode == 'RGBA':
        # Remove alpha channel
        return img.convert('RGB')

    if mode == 'P':
        # Palette -> RGB
        return img.convert('RGB')

    if mode == 'CMYK':
        return img.convert('RGB')

    if mode == 'LAB' or mode == 'HSV' or mode == 'YCbCr':
        # Convert to RGB (may lose color space information)
        return img.convert('RGB')

    # Default: try to convert to RGB
    return img.convert('RGB')


def to_single_channel(img: Image.Image, method: str = 'average') -> Image.Image:
    """
    Convert an image to single channel grayscale.

    Args:
        img: PIL Image object.
        method: How to convert multi-channel images:
            - 'average': average of channels (default)
            - 'luminosity': weighted average for RGB (0.299*R + 0.587*G + 0.114*B)
            - 'red', 'green', 'blue': take specific channel
            - 'max', 'min': take max or min across channels

    Returns:
        Single channel grayscale image (mode 'L').

    Raises:
        ValueError: If method is invalid.
    """
    if img.mode in ('L', '1', 'I', 'F'):
        # Already single channel
        if img.mode != 'L':
            return img.convert('L')
        return img

    # Convert to RGB if needed for channel operations
    if img.mode not in ('RGB', 'RGBA'):
        img = img.convert('RGB')

    # For multi-channel images
    if method == 'average':
        # Try to use numpy for accurate average if available
        if NUMPY_AVAILABLE:
            img_array = np.array(img)
            if len(img_array.shape) == 2:
                # Already single channel
                return Image.fromarray(img_array).convert('L')
            else:
                # Average across channels
                gray = np.mean(img_array, axis=2).astype(np.uint8)
                return Image.fromarray(gray).convert('L')
        else:
            # Fallback: use luminosity conversion (close to average)
            # Note: This is not mathematically the average, but a reasonable approximation
            # when numpy is not available
            return img.convert('L')

    elif method == 'luminosity':
        # Use PIL's built-in luminosity conversion
        return img.convert('L')

    elif method in ('red', 'green', 'blue'):
        # Extract specific channel
        channels = img.split()
        if len(channels) >= 3:
            channel_index = {'red': 0, 'green': 1, 'blue': 2}[method]
            return channels[channel_index].convert('L')
        else:
            # Not enough channels, use first channel
            return channels[0].convert('L')

    elif method in ('max', 'min'):
        # Take max or min across channels
        channels = img.split()
        if len(channels) == 0:
            return img.convert('L')

        # Convert all channels to 'L' mode for comparison
        channels_l = [chan.convert('L') for chan in channels]

        # Start with first channel
        result = channels_l[0]

        for chan in channels_l[1:]:
            if method == 'max':
                # Use lighter pixel (max)
                result = ImageChops.lighter(result, chan)
            else:  # 'min'
                # Use darker pixel (min)
                result = ImageChops.darker(result, chan)

        return result

    else:
        # Default PIL conversion
        return img.convert('L')


def to_rgb(img: Image.Image) -> Image.Image:
    """
    Convert an image to standard RGB format.

    This is similar to to_three_channel but ensures color space is RGB.

    Args:
        img: PIL Image object.

    Returns:
        RGB image.
    """
    return to_three_channel(img)


def to_grayscale(img: Image.Image, method: str = 'luminosity') -> Image.Image:
    """
    Convert an image to grayscale.

    Args:
        img: PIL Image object.
        method: Conversion method (see to_single_channel).

    Returns:
        Grayscale image (mode 'L').
    """
    return to_single_channel(img, method)


def fix_channels(img: Image.Image, target_channels: int) -> Image.Image:
    """
    Convert image to have specified number of channels.

    Args:
        img: PIL Image object.
        target_channels: Desired number of channels (1, 3, or 4).

    Returns:
        Image with target_channels channels.

    Raises:
        ValueError: If target_channels is not 1, 3, or 4.
    """
    if target_channels == 1:
        return to_single_channel(img)
    elif target_channels == 3:
        return to_three_channel(img)
    elif target_channels == 4:
        # Convert to RGBA
        if img.mode == 'RGBA':
            return img
        elif img.mode == 'RGB':
            return img.convert('RGBA')
        else:
            # Convert to RGB first, then add alpha
            rgb_img = to_three_channel(img)
            return rgb_img.convert('RGBA')
    else:
        raise ValueError(f"Unsupported target_channels: {target_channels}. Must be 1, 3, or 4.")


def get_dimensions(img: Image.Image) -> Tuple[int, int, int]:
    """
    Get image dimensions and channel count.

    Args:
        img: PIL Image object.

    Returns:
        Tuple of (width, height, channels).
    """
    width, height = img.size

    # Map PIL modes to channel count
    mode_to_channels = {
        '1': 1, 'L': 1, 'P': 1, 'I': 1, 'F': 1,
        'RGB': 3, 'YCbCr': 3, 'LAB': 3, 'HSV': 3,
        'RGBA': 4, 'CMYK': 4,
    }

    channels = mode_to_channels.get(img.mode, 1)
    return (width, height, channels)


def print_image_info(img: Image.Image, title: str = "Image Info") -> None:
    """
    Print detailed image information in a human-readable format.

    Args:
        img: PIL Image object.
        title: Optional title for the output.
    """
    width, height, channels = get_dimensions(img)
    info = detect_image_type(img)

    print(f"=== {title} ===")
    print(f"Dimensions: {width} x {height} pixels")
    print(f"Channels: {channels}")
    print(f"Mode: {img.mode}")
    print(f"Type: {info['type'].value}")
    print(f"Is grayscale: {info['is_actual_grayscale']}")

    # Show data type if available
    if img.mode in ('I', 'F'):
        dtype_map = {'I': '32-bit integer', 'F': '32-bit float'}
        print(f"Data type: {dtype_map.get(img.mode, 'unknown')}")
    else:
        print("Data type: 8-bit unsigned integer")

    print("=" * (len(title) + 8))